import numpy as np

class Evaluator(object):
    def __init__(self, num_class=2):
        """
        #2*2的矩阵,行代表ground truth类别(背景,建筑物),列代表preds的类别(背景,建筑物)

        """
        self.num_class = num_class
        self.confusion_matrix = np.zeros((self.num_class,)*2)

    def Pixel_Accuracy(self):
        '''
        正确的像素占总像素的比例PA
        '''
        Acc = np.diag(self.confusion_matrix).sum() / self.confusion_matrix.sum()
        return Acc

    def Mean_Intersection_over_Union(self):
        """
        MIoU计算
        """
        MIoU = np.diag(self.confusion_matrix) / (
                np.sum(self.confusion_matrix, axis=1) + np.sum(self.confusion_matrix, axis=0) -
                np.diag(self.confusion_matrix))
        MIoU = np.nanmean(MIoU)
        return MIoU

    def update_matrix(self, a, b):
        """
        混淆矩阵计算,设标签宽W，长H
        a：gt
        b：pred
        """
        n = self.num_class
        # a是转化成一维数组的标签，形状(H×W,)；b是转化成一维数组的标签，形状(H×W,)
        k = (a >= 0) & (a < n)
        # np.bincount计算了从0到(n**2-1)这n**2个数中每个数出现的次数，返回值形状(n, n)
        # 返回中，写对角线上的为分类正确的像素点
        self.confusion_matrix = self.confusion_matrix \
                                + np.bincount(n * a[k].astype(int) + b[k], minlength=n ** 2).reshape(n, n)

        return self.confusion_matrix

    def reset(self):
        self.confusion_matrix = np.zeros((self.num_class,)*2)


if __name__=="__main__":
    eval=Evaluator(2)

    gt_image=np.array([0,1,0,1,0,0,1,1,1])
    pr_image=np.array([0,1,0,0,1,1,1,0,0])

    matrix=eval.update_matrix(gt_image,pr_image)
    print(matrix)
    """
    [[2 2]
     [3 2]]
    """

    miou=eval.Mean_Intersection_over_Union()
    print(miou) # 0.2857142857142857

    PA=eval.Pixel_Accuracy()
    print(PA) # 0.4444444444444444